from deal_sgy import gain
import numpy as np
import copy

from util.plot import set_axis
from util.tool import get_index


# 绘制sgy数据
def draw_sgy(plt, sgy, deep_start = 20, deep_end = 300, time_start = 0, time_end = 1000):
    """
    绘制sgy数据
    :param plt: plt
    :param sgy: sgy数据对象
    :param deep_start: 起始深度，单位m，默认20
    :param deep_end: 结束深度，单位m，默认300
    :param time_start: 起始时间，单位ms，默认0
    :param time_end: 结束时间，单位ms，默认1000
    :return: ax, {
        "top": deep_start_idx,
        "bottom": deep_end_idx,
        "left": time_start_idx,
        "right": time_end_idx
    }
    """

    # ============================================================
    # 数据预处理
    # ============================================================

    data = copy.deepcopy(sgy["data"])  # 复制sgy数据避免对原数据产生影响

    time_start_idx = int(time_start / sgy["dt"])  # 计算时间点索引，开始
    time_end_idx = int(time_end / sgy["dt"])  # 计算时间点索引，结束
    deep_start_idx = get_index(sgy["cdp"], deep_start)
    deep_end_idx = get_index(sgy["cdp"], deep_end)
    sgy["deal_data"] = data[deep_start_idx:deep_end_idx:1, time_start_idx:time_end_idx:1]  # 对道和时间进行抽稀和截取，保存到deal_data属性中

    # sgy["deal_data"] = - sgy["deal_data"]  # 数据反极性

    if sgy["dt"] < 1:
        sca = int(1 / sgy["dt"])
        sgy["dt"] = 1.0
        sgy["deal_data"] = sgy["deal_data"][::, ::sca]  # 深度域不变，时间域抽稀
    print("==========*****==========\n数据抽稀完成")

    # 添加增益
    sgy["deal_data"] = gain(sgy, 500, "deal")
    print("==========*****==========\n数据增益完成")

    # ============================================================
    # 绘制数据
    # ============================================================
    amp_ratio = 0.55  # 振幅大小

    fig = plt.figure(facecolor='w', dpi=300)
    ax = fig.add_axes([0.07, 0.02, 0.9, 0.9])  # 拉平, left，bottom，width，height

    perc = 99.0
    rgba = [0, 0, 0, 1]
    sc = np.percentile(sgy["deal_data"], perc)  # Normalization factor，计算数据集中特定百分位数的值，相当于去除异常数据后的最大值
    wigdata = sgy["deal_data"]

    xpos = np.array(sgy["cdp"])  # 转为numpy数组，深度
    for brfjj, x, trace in zip(range(0, len(xpos) - 1), xpos, wigdata):
        print(f"第{brfjj}条绘制完成")
        amp = x - amp_ratio * 10 * trace / sc  # 每道数据的最大高度按10算
        t = np.arange(time_start, time_end, sgy["dt"])
        hypertime = np.linspace(t[0], t[-1], (10 * t.size - 1) + 1)
        hyperamp = np.interp(hypertime, t, amp)  # 差值后的道数据
        ax.plot( hypertime, hyperamp,'k', lw=1.2 * 0.2)
        ax.fill_between(hypertime, hyperamp, x, where=hyperamp<x, facecolor=rgba, lw=0, )

    # ============================================================
    # 设置坐标轴
    # ============================================================

    x_min = time_start
    x_max = time_end
    y_min = sgy["cdp"][deep_start_idx]
    y_max = sgy["cdp"][deep_end_idx]

    set_axis(ax, fig, plt, x_min, x_max, y_min, y_max, "Time(ms)", "Deep(m)")

    return ax, {
        "top": deep_start_idx,
        "bottom": deep_end_idx,
        "left": time_start_idx,
        "right": time_end_idx
    }